import os
from model.aoi import Lot, Wafer, Image, Failure
from model.pmi import PMILot, PMIWafer, PMIImage, PMIDevice
import json
from utils.wafer import get_wafer_map
from utils.file import rename_pmi
from datetime import datetime, timedelta
from utils.database import SessionLocal
from utils.msg_redis import MessageRedis
from decimal import Decimal
import time
import sys
import asyncio
from weice_package import weice
# sys.path.append("/data/project/weice_gpu/gpu_e2e/weice_package")

# import weice

redis = MessageRedis()


def detect_aoi_lot(det_model, seg_model, cls_model, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale, manager):

    # 后台任务只能启动一次
    # redis.set_string("task", "start")

    with SessionLocal() as session:
    
        try:
            for folder in os.listdir(lot_folder):
                
                # 获取 lot 文件夹
                lotname = folder
                lot_dir = os.path.join(lot_folder, folder)

                lot = session.query(Lot).filter_by(lot_name=lotname).first()

                if lot:
                    print("1111111111111已存在数据库中: ", lotname)
                    continue
                else:
                    lot = Lot(lot_name=lotname, dir_path=lot_dir)
                    print("222222222222222222开始处理lot: ", lotname)
                    session.add(lot)
                    session.commit()

                    # 开始推理每个 wafer folder
                    for file in os.listdir(lot_dir):
                        # 获取 wafer 矩阵文件
                        if file.endswith('.txt'):
                            wafer_index = file.split("-")[1]
                            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                            json_data = {"msg": 'Lot %s Wafer %s 开始处理' % (lotname, wafer_index), "time": current_time, "status": "start", "level": "info"}
                            asyncio.run(manager.send_json("aoi", json_data))

                            wafer_folder = os.path.join(lot_dir, wafer_index) + "/"

                            wafer_txt = os.path.join(lot_dir, file)
                            print("aaaaaaaaaaaaaaaaaa开始处理wafer: ", wafer_index)

                            t1 = time.time()

                            predict_data = weice.run_aoi(det_model, seg_model, cls_model, wafer_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale)

                            t2 = time.time()
                            print("bbbbbbbbbbbbbbbbbbbb结束处理wafer: ", wafer_index)

                            print("推理耗时：", t2-t1)

                            # 生成 wafer map 矩形图
                            wafer_data = get_wafer_map(wafer_txt, wafer_folder, predict_data)

                            wafer = Wafer(wafer_id=wafer_index, lot_code=lot.lot_code, wafer_data=json.dumps(wafer_data))
                            session.add(wafer)
                            session.commit()

                            # 每张图片的推理结果
                            for image_predict in predict_data:
                                failure_id = 1
                                detection_info = {}
                                segmentation_info = {}
                                image_name = image_predict["filename"]

                                wuran_box = image_predict.get("wuran_box", [])
                                wuran_crop = image_predict.get("wuran_crop", [])

                                zhenhen_box = image_predict.get("zhenhen_box", [])
                                zhenhen_crop = image_predict.get("zhenhen_crop", [])
                                # zhenhen_num_error = image_predict.get("zhenhen_num_error", [])

                                can_wuran_box = image_predict.get("can_wuran_box", [])

                                madian_box = image_predict.get("madian_box", [])

                                # if zhenhen_num_error:
                                #     detection_info[4] = zhenhen_num_error
                                #     failure_id = 4

                                if zhenhen_box:
                                    detection_info[3] = zhenhen_box
                                    failure_id = 3
                                
                                if zhenhen_crop:
                                    segmentation_info[3] = zhenhen_crop
                                    failure_id = 3

                                if can_wuran_box:
                                    detection_info[5] = can_wuran_box
                                    failure_id = 5

                                if wuran_box:
                                    detection_info[2] = wuran_box
                                    failure_id = 2

                                if wuran_crop:
                                    segmentation_info[2] = wuran_crop
                                    failure_id = 2

                                if madian_box:
                                    detection_info[6] = madian_box
                                    failure_id = 6

                                file_path = os.path.join(lot_dir, wafer_index, image_name)

                                if detection_info:
                                    detection_info = json.dumps(detection_info)
                                else:
                                    detection_info = ""

                                if segmentation_info:
                                    segmentation_info=json.dumps(segmentation_info)
                                else:
                                    segmentation_info = ""

                                image = Image(
                                    wafer_code = wafer.wafer_code, 
                                    file_path = file_path,
                                    file_name = image_name,
                                    detection_info=detection_info,
                                    segmentation_info=segmentation_info,
                                    failure_id = failure_id
                                )
                                
                                session.add(image)

                            # wafer 处理完
                            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                            json_data = {"msg": 'Lot %s Wafer %s 处理完成' % (lotname, wafer_index), "time": current_time, "status": "end", "level": "info"}  
                            asyncio.run(manager.send_json("aoi", json_data))

                # lot 处理完
                current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                json_data = {"msg": 'Lot %s 处理完成' % lotname, "time": current_time, "status": "completed", "level": "info"}  

                lot.status = "COMPLETED"
                session.commit()

                asyncio.run(manager.send_json("aoi", json_data))
    

        except Exception as e:
            print("Exception: ", e)
            current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
            json_data = {"msg": 'Lot %s 处理出错' % lotname, "time": current_time, "level": "error"}
            asyncio.run(manager.send_json("aoi", json_data))
            session.rollback()

        # 后台任务执行完成
        # redis.set_string("task", "end")


# 根据 pmi_num 区别不同机台任务
def detect_pmi_folder(pmi_yolox_model, pmi_seg_model, pmi_num, pmi_folder, distance, set_per_area, area_scale, manager):

    # 后台任务只能启动一次
    # redis.set_string("task", "start")
      
    with SessionLocal() as session:
        try:
            pmi_device = session.query(PMIDevice).filter_by(device_id=pmi_num).first()
            device_code = pmi_device.device_code
            
            # 获取 pmi 机台编号文件夹
            lot_folder = os.path.join(pmi_folder, pmi_num, "pmi_images", "Failed")

            for lotname in os.listdir(lot_folder):
                lot_dir = os.path.join(lot_folder, lotname)

                # 判断是否是一个新的 lot，如果是的，需要写入数据库
                lot = session.query(PMILot).filter_by(lot_name=lotname, device_code=device_code).first()

                if not lot:
                    print("00000000不存在 pmi 数据库中: ", lotname)
                    lot = PMILot(lot_name=lotname, dir_path=lot_dir, device_code=device_code)
                    session.add(lot)
                    session.commit()
                else:
                    print("1111111111111已存在 pmi 数据库中: ", lotname)
                lot_code = lot.lot_code

                print("222222222222222222开始处理pmi lot: ", lotname)

                for waferid in os.listdir(lot_dir):
                    if_pass = True
                    wafer = session.query(PMIWafer).filter_by(lot_code=lot_code, wafer_id=waferid).first()
                    if not wafer:
                        wafer = PMIWafer(wafer_id=waferid, lot_code=lot.lot_code)
                        session.add(wafer)
                        session.commit()

                    wafer_code = wafer.wafer_code
                    wafer_dir = os.path.join(lot_dir, waferid)  + "/"

                    img_files = []

                    # 每次推理一个 wafer 文件夹里面的图片
                    for img_name in os.listdir(wafer_dir):
                        image = session.query(PMIImage).filter_by(wafer_code=wafer_code, file_name=img_name).first()
                        if not image:
                            # 重命名 pmi 图片名
                            new_image_name = rename_pmi(wafer_dir, img_name)

                            img_files.append(new_image_name)

                    # 如果 wafer 目录中没有新增图片，则继续遍历其它 wafer 目录
                    if not img_files:
                        continue
                    
                    current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')

                    json_data = {"msg": 'Lot %s Wafer %s 正在处理' % (lotname, waferid), "time": current_time, "status": "start", "level": "info"}
                    # redis.set_message("pmi", json_data)
                    print("aaaaaaaaaaaaaaaaaa开始处理pmi wafer: ", waferid)
                    asyncio.run(manager.send_json("pmi", json_data))

                    predict_data = weice.run_pmi(pmi_yolox_model,pmi_seg_model,wafer_dir,img_files,distance,set_per_area,area_scale)

                    print("bbbbbbbbbbbbbbbbbbbb结束处理pmi wafer: ", waferid)

                    for image_predict in predict_data:
                        failure_id = 1
                        detection_info = {}
                        segmentation_info = {}
                        image_name = image_predict["filename"]

                        image_path = os.path.join(wafer_dir, image_name)

                        if image_name == "F2..61..1.19.jpg":
                            print(image_predict)

                        wuran_box = image_predict.get("wuran_box", [])
                        wuran_crop = image_predict.get("wuran_crop", [])
                        zhenhen_box = image_predict.get("zhenhen_box", [])
                        zhenhen_crop = image_predict.get("zhenhen_crop", [])
                        zhenhen_num_error = image_predict.get("zhenhen_num_error", [])

                        # if wuran_box:
                        #     detection_info[2] = wuran_box
                        if zhenhen_box:
                            detection_info[3] = zhenhen_box
                            failure_id = 3
                        # if zhenhen_num_error:
                        #     detection_info[4] = zhenhen_num_error

                        # if wuran_crop:
                        #     segmentation_info[2] =  wuran_crop
                        if zhenhen_crop:
                            segmentation_info[3] = zhenhen_crop
                            failure_id = 3

                        if detection_info:
                            if_pass = False
                            detection_info = json.dumps(detection_info)
                        else:
                            detection_info = ""

                        if segmentation_info:
                            if_pass = False
                            segmentation_info=json.dumps(segmentation_info)
                        else:
                            segmentation_info = ""

                        image = PMIImage(
                            wafer_code=wafer_code, file_name=image_name, file_path=image_path,
                            detection_info=detection_info,
                            segmentation_info=segmentation_info,
                            failure_id = failure_id)
                        session.add(image)
                    
                    current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                    json_data = {"msg": 'Lot %s Wafer %s 处理结束' % (lotname, waferid), "time": current_time, "status": "end", "level": "info"}           
                    # redis.set_message("pmi", json_data)
                    asyncio.run(manager.send_json("pmi", json_data))

                    session.commit()

                    # 如果检测出错误图片
                    if not if_pass:
                        json_data = {"msg": 'Lot %s Wafer %s NO PASS' % (lotname, waferid), "time": current_time, "status": "end", "level": "error", "wafer_code": wafer_code}     
                    else:
                        json_data = {"msg": 'Lot %s Wafer %s PASS' % (lotname, waferid), "time": current_time, "status": "end", "level": "info"}      
                    # redis.set_message("pmi", json_data)
                    asyncio.run(manager.send_json("pmi", json_data))
                    
                # lot 处理完
                current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
                json_data = {"msg": 'Lot %s 处理完成' % lotname, "time": current_time, "status": "completed"} 

                lot.status = "COMPLETED"
                session.commit()
                asyncio.run(manager.send_json("pmi", json_data))

        except Exception as e:
            print("Exception: ", e)
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            json_data = {"msg": '处理出错: %s' % str(e), "time": current_time, "level": "error"}
            redis.set_message("pmi", json_data)
            # session.rollback()

        # 后台任务执行完成
        # redis.set_string("task", "end")


def detect_aoi_wafer(aoi_yolox_model, aoi_seg_model, cls_model, wafer_folder, distance=6, set_num_zhenhen=3, set_diagonal_length=30, set_area=25, set_per_area=0.25, area_scale=0):

    wafer_folder = wafer_folder + "/"
      
    predict_data = weice.run_aoi(aoi_yolox_model, aoi_seg_model, cls_model, wafer_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale)

    ng_predict_data = []
    good_predict_data = []
    all_failure_num = {}

    # 每张图片的推理结果
    for image_predict in predict_data:
        item = {}
        failure_id = 1
        detection_info = {}
        segmentation_info = {}
        image_name = image_predict["filename"]

        item["file_name"] = image_name
        # 所有错误信息，为了计算出正常的图片
        item["detection_info"] = []
        item["segmentation_info"] = []

        wuran_box = image_predict.get("wuran_box", [])
        wuran_crop = image_predict.get("wuran_crop", [])

        can_wuran_box = image_predict.get("can_wuran_box", [])

        zhenhen_box = image_predict.get("zhenhen_box", [])
        zhenhen_crop = image_predict.get("zhenhen_crop", [])
        # zhenhen_num_error = image_predict.get("zhenhen_num_error", [])

        madian_box = image_predict.get("madian_box", [])

        # if zhenhen_num_error:
        #     detection_info[4] = zhenhen_num_error
        #     failure_id = 4

        if zhenhen_box:
            detection_info[3] = zhenhen_box
            failure_id = 3
        
        if zhenhen_crop:
            segmentation_info[3] = zhenhen_crop
            failure_id = 3

        if can_wuran_box:
            detection_info[5] = can_wuran_box
            failure_id = 5

        if wuran_box:
            detection_info[2] = wuran_box
            failure_id = 2

        if wuran_crop:
            segmentation_info[2] = wuran_crop
            failure_id = 2

        if madian_box:
            detection_info[6] = madian_box
            failure_id = 6

        if failure_id == 1:
            good_predict_data.append(item) 
            continue

        if detection_info:
            for key, value in detection_info.items():

                for detection in value:
                    bbox = detection[0]   
                    distance = None   
                    if str(key) == "3":
                        distance = detection[2]
                    item["detection_info"].append([key, bbox, None, distance]) 

        if segmentation_info:
            for key, value in segmentation_info.items():

                # 每个 pad 的 segmentation 在一起
                for segmentation in value:
                    points = segmentation[0]
                    all_xy = []
                    bbox = segmentation[2]
                    area_rate = None
                    if str(key)=="3":
                        area_rate = int(Decimal(segmentation[3]).quantize(Decimal("0.01"), rounding="ROUND_HALF_UP")*100)
                    # points 指的是一张图片的所有轮廓
                    for point in points:
                        # point 指一个 pad 的轮廓
                        point_xy = []
                        for xy in point:
                            x = xy[0][0] + bbox[0]
                            y = xy[0][1] + bbox[1]
                            point_xy.append([x, y])
                        all_xy.append(point_xy)
                    item["segmentation_info"].append([key, bbox, all_xy, area_rate])  

        ng_predict_data.append(item) 

        # if all_failure_num.get(failure_id):
        #     all_failure_num[failure_id] += 1
        # else:
        #     all_failure_num[failure_id] = 1

    class_label = []

    with SessionLocal() as session:
        all_failures = session.query(Failure).filter_by(del_flag="0")
        for failure in all_failures:
            label_item = {}
            label_item["failure_name"] = failure.failure_name
            label_item["failure_id"] = failure.failure_id
            class_label.append(label_item)

    return {"predict_data": good_predict_data, "ng_predict_data": ng_predict_data, "class_label": class_label}


def delete_aoi_db():
    current_time = (datetime.utcnow() + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
    json_data = {"msg": '清空 aoi 数据库', "time": current_time, "status": "delete", "level": "info"}
    # redis.set_message("aoi", json_data)